#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：量化交易探究 
@File ：连续交易模拟.py
@IDE  ：PyCharm 
@Author ：李涵彬
@Date ：2025/3/16 上午8:10 
"""

import random
import numpy as np
import matplotlib

matplotlib.use('TkAgg')
import matplotlib.pyplot as plt


class TraderSimulator:
	def __init__(
			self,
			initial_balance: float,
			risk_per_trade: float,
			win_rate: float,
			reward_ratio_distribution: tuple,
			risk_free_rate: float = 0.02,
			trading_days_per_year: int = 252
	):
		self.initial_balance = initial_balance
		self.risk_per_trade = risk_per_trade
		self.win_rate = win_rate
		self.reward_mean, self.reward_std = reward_ratio_distribution
		self.risk_free_rate = risk_free_rate
		self.trading_days = trading_days_per_year

	def simulate_trades(self, num_trades: int) -> dict:
		balance = self.initial_balance
		balances = [balance]
		results = []
		profits = []
		losses = []
		daily_returns = []

		for _ in range(num_trades):
			is_win = random.random() < self.win_rate
			results.append(is_win)

			# 动态生成盈亏比
			reward_ratio = max(0.1, np.random.normal(self.reward_mean, self.reward_std))

			risk_amount = balance * self.risk_per_trade

			if is_win:
				profit = risk_amount * reward_ratio
				balance += profit
				profits.append(profit)
			else:
				loss = -risk_amount
				balance += loss
				losses.append(loss)

			# 计算日收益率（假设每个交易日进行一次交易）
			daily_return = (balance - balances[-1]) / balances[-1] if balances[-1] != 0 else 0
			daily_returns.append(daily_return)
			balances.append(balance)

		return {
			"balances": balances,
			"results": results,
			"profits": profits,
			"losses": losses,
			"daily_returns": daily_returns
		}

	def calculate_sharpe_ratio(self, daily_returns: list, balances: list) -> float:
		annual_return = (self.initial_balance / (balances[-1] if balances[-1] != 0 else 1)) ** self.trading_days - 1
		annualized_return = (annual_return + 1) ** (1 / self.trading_days) - 1  # 年化收益率
		volatility = np.std(daily_returns) * np.sqrt(self.trading_days)
		sharpe = (annualized_return - self.risk_free_rate) / volatility if volatility != 0 else 0
		return sharpe

	def calculate_max_drawdown(self, balances: list) -> tuple:
		max_dd_abs = 0.0
		max_dd_pct = 0.0
		max_dd_start = 0
		max_dd_end = 0
		current_peak = balances[0]
		current_start = 0  # 显式初始化 current_start（关键修复）

		for i in range(1, len(balances)):
			current = balances[i]
			if current > current_peak:
				current_peak = current
				current_start = i
			else:
				# 计算回撤...
				current_dd_abs = current_peak - current
				current_dd_pct = (current_peak - current) / current_peak

				if current_dd_abs > max_dd_abs:
					max_dd_abs = current_dd_abs
					max_dd_pct = current_dd_pct
					max_dd_start = current_start  # 使用 current_start
					max_dd_end = i
		return max_dd_abs, max_dd_pct, max_dd_start, max_dd_end

	def calculate_consecutive_counts(self, results: list) -> tuple:
		consecutive_win = 0
		max_win = 0
		consecutive_loss = 0
		max_loss = 0

		for r in results:
			if r:
				consecutive_win += 1
				consecutive_loss = 0
				if consecutive_win > max_win:
					max_win = consecutive_win
			else:
				consecutive_loss += 1
				consecutive_win = 0
				if consecutive_loss > max_loss:
					max_loss = consecutive_loss

		return max_win, max_loss

	def calculate_profit_metrics(self, profits: list, losses: list) -> dict:
		total_profit = sum(profits) if profits else 0
		total_loss = -sum(losses) if losses else 0
		avg_profit = total_profit / len(profits) if profits else 0
		avg_loss = total_loss / len(losses) if losses else 0
		profit_factor = total_profit / total_loss if total_loss != 0 else np.inf

		return {
			"total_profit": total_profit,
			"total_loss": total_loss,
			"avg_profit": avg_profit,
			"avg_loss": avg_loss,
			"profit_factor": profit_factor
		}

	def calculate_risk_metrics(self, balances: list, daily_returns: list) -> tuple:
		volatility = np.std(daily_returns) * np.sqrt(self.trading_days) if daily_returns else 0
		max_dd_abs, max_dd_pct, _, _ = self.calculate_max_drawdown(balances)
		return volatility, max_dd_abs, max_dd_pct

	def plot_funds_curve(self, balances: list, max_dd_start: int, max_dd_end: int):
		plt.figure(figsize=(14, 7))
		plt.plot(balances, label='Account Balance')
		plt.title('Funds Curve with Risk Metrics')
		plt.xlabel('Trade Number')
		plt.ylabel('Balance')
		plt.grid(True)

		# 标记最大回撤区间
		plt.axvspan(max_dd_start, max_dd_end, alpha=0.2, color='red', label='Max Drawdown')
		plt.legend()
		plt.show()

	def run_simulation(self, num_trades: int) -> dict:
		data = self.simulate_trades(num_trades)
		balances = data["balances"]
		results = data["results"]
		profits = data["profits"]
		losses = data["losses"]
		daily_returns = data["daily_returns"]

		# 计算关键指标
		max_win, max_loss = self.calculate_consecutive_counts(results)
		profit_metrics = self.calculate_profit_metrics(profits, losses)
		volatility, max_dd_abs, max_dd_pct = self.calculate_risk_metrics(balances, daily_returns)
		sharpe = self.calculate_sharpe_ratio(daily_returns, balances)

		return {
			"final_balance": balances[-1],
			"total_trades": num_trades,
			"win_rate": sum(results) / num_trades if num_trades else 0,
			"sharpe_ratio": sharpe,
			"volatility": volatility,
			"max_drawdown_abs": max_dd_abs,
			"max_drawdown_pct": max_dd_pct,
			"profit_factor": profit_metrics["profit_factor"],
			"avg_profit": profit_metrics["avg_profit"],
			"avg_loss": profit_metrics["avg_loss"],
			"max_consecutive_win": max_win,
			"max_consecutive_loss": max_loss
		}


def main():
	config = {
		'initial_balance': 100000,
		'risk_per_trade': 0.02,  # 2%风险
		'win_rate': 0.55,  # 55%胜率
		'reward_distribution': (2.0, 0.5),  # 盈亏比正态分布（均值2.0，标准差0.5）
		'num_trades': 250  # 250个交易日
	}

	simulator = TraderSimulator(
		config['initial_balance'],
		config['risk_per_trade'],
		config['win_rate'],
		config['reward_distribution']
	)

	results = simulator.run_simulation(config['num_trades'])

	print("=== 交易模拟结果 ===")
	print(f"最终账户余额: {results['final_balance']:.2f} 元")
	print(f"总交易次数: {results['total_trades']}")
	print(f"胜率: {results['win_rate'] * 100:.1f}%")
	print(f"夏普比率: {results['sharpe_ratio']:.2f}")
	print(f"年化波动率: {results['volatility'] * 100:.2f}%")
	print(f"最大回撤金额: {results['max_drawdown_abs']:.2f} 元 ({results['max_drawdown_pct'] * 100:.1f}%)")
	print(f"盈利因子: {results['profit_factor']:.2f}")
	print(f"平均盈利: {results['avg_profit']:.2f} 元")
	print(f"平均亏损: {results['avg_loss']:.2f} 元")
	print(f"最大连续盈利次数: {results['max_consecutive_win']}")
	print(f"最大连续亏损次数: {results['max_consecutive_loss']}")

	# 绘制资金曲线（需手动调用）
	sim_data = simulator.simulate_trades(config['num_trades'])
	_, _, max_dd_start, max_dd_end = simulator.calculate_max_drawdown(sim_data["balances"])
	simulator.plot_funds_curve(sim_data["balances"], max_dd_start, max_dd_end)


if __name__ == "__main__":
	main()
